#!/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================

import os
import torch
import numpy as np

def token_shift(x: torch.Tensor, rwkvag: torch.Tensor, h0: torch.Tensor) -> torch.Tensor:
    """
    通道混合函数。

    Args:
        x (torch.Tensor): 输入张量，形状为[Batch, Seq_length, N_embd]。
        rwkvag (torch.Tensor): 输入张量(模型权重），形状为[6, 1, N_embd]。
        h0 (torch.Tensor): 输入张量(state状态)，形状为[Batch, 1, N_embd]。
    Returns:
        xr, xw, xk, xv, xa, xg (torch.Tensor):输出r, w, k, v, a, g混合后的张量,形状均为[Batch, Seq_length, N_embd]
        ht (torch.Tensor): 输出张量(state状态),形状为[Batch, 1, N_embd]。
    """
    batch_size, seq_length = x.shape[0], x.shape[1]
    if seq_length == 1:
        sx = (h0 - x)
        ht = x
    else:
        h0 = h0.view(batch_size,1,-1)
        xx = torch.cat([h0, x[:, :-1, :]], dim=1)
        sx = (xx - x)
        ht = x[:, -1, :]
    xr = x + rwkvag[0] * sx
    xw = x + rwkvag[1] * sx
    xk = x + rwkvag[2] * sx
    xv = x + rwkvag[3] * sx
    xa = x + rwkvag[4] * sx
    xg = x + rwkvag[5] * sx
    return xr, xw, xk, xv, xa, xg, ht

def gen_golden_data_simple():
    dtype = np.float16

    B = 77
    T = 77
    C = 2560

    x_np = np.random.uniform(-1, 1, [B, T, C]).astype(dtype)
    rwkvag_np = np.random.uniform(-1, 1, [6, 1, 1, C]).astype(dtype)
    h0_np = np.random.uniform(-1, 1, [B, 1, C]).astype(dtype)

    x = torch.from_numpy(x_np)
    rwkvag = torch.from_numpy(rwkvag_np)
    h0 = torch.from_numpy(h0_np)

    xr, xw, xk, xv, xa, xg, ht = token_shift(x, rwkvag, h0)

    os.system("mkdir -p input")
    x_np.astype(dtype).tofile("./input/input_x.bin")
    rwkvag_np.astype(dtype).tofile("./input/input_rwkvag.bin")
    h0_np.astype(dtype).tofile("./input/input_h0.bin")

    os.system("mkdir -p output")
    xr.cpu().numpy().tofile("./output/out_xr.bin")
    xw.cpu().numpy().tofile("./output/out_xw.bin")
    xk.cpu().numpy().tofile("./output/out_xk.bin")
    xv.cpu().numpy().tofile("./output/out_xv.bin")
    xa.cpu().numpy().tofile("./output/out_xa.bin")
    xg.cpu().numpy().tofile("./output/out_xg.bin")
    ht.cpu().numpy().tofile("./output/out_ht.bin")

if __name__ == "__main__":
    gen_golden_data_simple()

